# -*- coding: utf-8 -*-
"""
PBAR_Zspec_Compare1st2ndSets.py
Created on Thu Apr 03 15:01:03 2014

@author: jkwong
"""

import os, cPickle
import numpy as np
import matplotlib.pyplot as plt


# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2

datAll = []

#Load the two sets
basepath = r'C:\Users\jkwong\Documents\Work\PBAR\data'
fullFilename = os.path.join(basepath, 'zspec1stSet.dat')
with open(fullFilename ,'rb') as fid:
    print('reading %s' %fullFilename)
    datAll.append(cPickle.load(fid))

basepath = r'C:\Users\jkwong\Documents\Work\PBAR\data4\Mar-files'
fullFilename = os.path.join(basepath, 'zspec2ndSet.dat')
with open(fullFilename ,'rb') as fid:
    print('reading %s' %fullFilename)
    datAll.append(cPickle.load(fid))

# plot the features vs counts

featureIndex = 1
plt.figure()
plt.grid()
for i in xrange(len(datAll)):
    for j in xrange(2):
        if i == 0:
            markerSize = 3
            if j == 0:
                plotColor = 'r'
            else:
                plotColor = 'b'
        elif i == 1:
            markerSize = 6            
            if j == 0:
                plotColor = 'm'
            else:
                plotColor = 'c'

        cut = datAll[i]['statsMatrix'][:,-1] == j
        plt.plot(datAll[i]['countMatrix'][cut], datAll[i]['statsMatrix'][cut,featureIndex], '.r', \
            color = plotColor, marker = markerTypes[i+1], markersize = markerSize,  alpha = 0.2, \
            label = 'Set%d, %s' %(i+1, datAll[i]['groupNamesExportList'][j]) )

plt.legend()
plt.xlabel('Counts', fontsize = 16)
plt.ylabel('Mean Bin', fontsize = 16)


# plot the features vs counts

featureIndex = 1
featureIndex = np.where(statsCollapseListALL == 'binSkew')[0][0]
for i in xrange(len(datAll)):
    plt.figure()
    plt.grid()

    for j in xrange(2):
        cut = datAll[i]['statsMatrix'][:,-1] == j
        plt.plot(datAll[i]['countMatrix'][cut], datAll[i]['statsMatrix'][cut,featureIndex], '.r', \
            color = plotColors[j], marker = markerTypes[i+1], markersize = markerSize,  alpha = 0.2, \
            label = 'Set%d, %s' %(i+1, datAll[i]['groupNamesExportList'][j]) )
    plt.legend()
    plt.xlabel('Counts', fontsize = 16)
    plt.ylabel('Mean Bin', fontsize = 16)
    plt.title('Set %d' %(i+1))
    plt.axis((0,400,0,70))
    
    if savePlots:
        filename = os.path.join(plotPath, 'MeanBin_Count_Set%d.png' %(i+1))
        plt.savefig(filename)

# plot of features vs counts
# pb, not pb


featureIndex = np.where(statsCollapseListALL == 'binMean_g1')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binMean')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binSkew')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binKur')[0][0]
#featureIndex = 12

plt.figure()
plt.grid()

# all lead, all no lead
cut = statsMatrix[:,-1] == 0
plt.plot(countMatrix[cut], statsMatrix[cut,featureIndex], 'sb', markersize = 7, alpha = 0.2, label = 'Not Pb, Set %d' %setNum)
cut = statsMatrix[:,-1] == 1
plt.plot(countMatrix[cut], statsMatrix[cut,featureIndex], 'sr', markersize = 7, alpha = 0.2, label = 'Pb, Set %d' %setNum)
#
# #set 2a
cut = statsMatrix_0[:,-1] == 0
plt.plot(countMatrix_0[cut], statsMatrix_0[cut,featureIndex], '^b', markersize = 7, alpha = 0.3, label = 'Not Pb, Set 2a')
cut = statsMatrix_0[:,-1] == 1
plt.plot(countMatrix_0[cut], statsMatrix_0[cut,featureIndex], '^r', markersize = 7, alpha = 0.3, label = 'Pb, Set 2a')
# set 2b
cut = statsMatrix_1[:,-1] == 0
plt.plot(countMatrix_1[cut], statsMatrix_1[cut,featureIndex], 'dc', markersize = 7, alpha = 0.7, label = 'Not Pb, Set 2b')
cut = statsMatrix_1[:,-1] == 1
plt.plot(countMatrix_1[cut], statsMatrix_1[cut,featureIndex], 'dm', markersize = 7, alpha = 0.7, label = 'Pb, Set 2b')

#
#detectorList = goodDetectorsList[(goodDetectorsList > 40) & (goodDetectorsList < 120)]
##groupName = 'PbNOT_1'
##groupName = 'Al_1'
#groupName = 'Fe_0'
#
#for i, index in enumerate(datasetGroupsIndices[groupName]):
#    plt.plot(stats['count'][index,detectorList], stats[statsCollapseListALL[featureIndex]][index,detectorList], \
#        'xm', color = plotColors[i%7], marker = markerTypes[i/7], linestyle = '-', \
#        markersize = 7, alpha = 0.7, label = '%s, %s'%(filenameList[index], datasetDescription[index]))
#
#groupName = 'Al_0'
#for i, index in enumerate(datasetGroupsIndices[groupName]):
#    plt.plot(stats['count'][index,detectorList], stats[statsCollapseListALL[featureIndex]][index,detectorList], \
#        'xm', color = plotColors[i%7], marker = markerTypes[(i)/7 + 2], linestyle = '-', \
#        markersize = 7, alpha = 0.7, label = '%s, %s'%(filenameList[index], datasetDescription[index]))


plt.xlabel('Count')
plt.ylabel(statsCollapseListALL[featureIndex])
plt.legend(loc =1)

if featureIndex == 9:
    axis((0, 400, 0, 3e4))
elif featureIndex == 11:
    axis((0, 400, 0, 2))
elif featureIndex == 12:
    axis((0, 400, 2, 7))
    
axis((0, 400, 3, 7))
axis((0, 400, 10, 50))


# Make plot so of features vs counts one at a time and save to file

for featureIndex in xrange(len(statsCollapseListALL)):
    plt.figure()
    plt.grid()
    
    cut = statsMatrix[:,-1] == 0
    plt.plot(countMatrix[cut], statsMatrix[cut,featureIndex], '.k', alpha = 0.5)
    
    cut = statsMatrix[:,-1] == 1
    plt.plot(countMatrix[cut], statsMatrix[cut,featureIndex], '.r', alpha = 0.5)

    temp = statsMatrix[cut,featureIndex]
    cut2 = (countMatrix[cut] > 100) & (countMatrix[cut] < 300)

    axis((0, 400, 0, 2*temp[cut2].mean()))
    
    xlabel('Count')
    ylabel(statsCollapseListALL[featureIndex])
    if savePlots:
        filename = os.path.join(plotPath, '%s_Count_Set2_Interp.png' %statsCollapseListALL[featureIndex])
        plt.savefig(filename)


# discriminatino performance plots

# get all the values from the  list dict and put into nice arrays

#plot f1

scores = []

#scoresMean[metric name][classification algorithm][set number] =
#scoresSTD[[metric name][classification algorithm][set number] =

classifierKeys = scoresTrialAllAll[0][0].keys()
maximumFeatures = len(scoresTrialAllAll[0])


scoresMean = []
scoresSTD = []
for setIndex, setNumber in enumerate([1, 2, 3, 4]):
    scoresMean.append({})    
    scoresSTD.append({})
    for key in classifierKeys: # cycle through algorithms
        scoresMean[setIndex][key] = {}
        scoresSTD[setIndex][key] = {}
        for statNameIndex, statName in enumerate(['p', 'r', 'f', 's']):
            scoresMean[setIndex][key][statName] = np.zeros((len(scoresTrialAllAll), maximumFeatures))
            scoresSTD[setIndex][key][statName] = np.zeros((len(scoresTrialAllAll), maximumFeatures))

for setIndex, setNumber in enumerate([1, 2, 3, 4]):
    fullFilename = os.path.join(basepath, 'zspec%dSet.dat' %setNumber)
    with open(fullFilename ,'rb') as fid:
        datTemp = cPickle.load(fid)
    scoresTrialAllAll = dat['scoresTrialAllAll']
    for key in classifierKeys: # cycle through algorithms
        for c in xrange(len(scoresTrialAllAll)): # cycle through count windows
            temp = np.zeros((scoresTrialAllAll[0][0]['lr'].shape[0], scoresTrialAllAll[0][0]['lr'].shape[1], maximumFeatures))
            for numFeaturesKeepIndex, numFeaturesKeep in enumerate(xrange(1, maximumFeatures+1)):
                temp[:,:,numFeaturesKeepIndex] = scoresTrialAllAll[c][numFeaturesKeepIndex][key]
            tempMean = temp.mean(0)
            tempSTD = temp.std(0)
            # means
            scoresMean[setIndex][key]['p'][c,:] = tempMean[0,:]
            scoresMean[setIndex][key]['r'][c,:] = tempMean[1,:]
            scoresMean[setIndex][key]['f'][c,:] = tempMean[2,:]
            scoresMean[setIndex][key]['s'][c,:] = tempMean[4,:]
            # standard deviation
            scoresSTD[setIndex][key]['p'][c,:] = tempSTD[0,:]
            scoresSTD[setIndex][key]['r'][c,:] = tempSTD[1,:]
            scoresSTD[setIndex][key]['f'][c,:] = tempSTD[2,:]
            scoresSTD[setIndex][key]['s'][c,:] = tempSTD[4,:]
    del datTemp
            
countRangesCenterList = np.array(  [mean(countRange) for countRange in countRangesList]   )

setNameList = ['1', '2', '2a', '2b']
metricName = 'f'
plt.figure()
plt.grid()
for setIndex, setNumber in enumerate([1, 2, 3, 4]):
    for key in classifierKeys: # cycle through algorithms
        if key == 'lda':
            linestyle = lineStyles[0]
        elif key == 'lr':
            linestyle = lineStyles[1]
        x = scoresMean[setIndex][key][metricName][:,5]
#        x = scoresMean[setIndex][key][metricName].max(1)
        x_mean = x[1:8].mean()
        plt.plot(countRangesCenterList, scoresMean[setIndex][key][metricName][:,5], \
            linestyle = linestyle, color = plotColors[setIndex%7], label = 'Set %s, %s, <F1> = %3.3f' %(setNameList[setIndex], key, x_mean))
plt.xlabel('Count Rate (Normalized to 60 Hz)', fontsize = 16)
plt.ylabel('F1', fontsize = 16)
plt.legend(loc = 3)
plt.axis((0, 350, 0.8, 1.0))

# scoresTrial[key][numTrial,:] = np.array([p, r, fbeta_score, support,  sum(y_test == pred) / float(len(pred))])


# Monthly report


#featureIndex = np.where(statsCollapseListALL == 'binMean_g1')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binMean')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binKur')[0][0]
featureIndex = np.where(statsCollapseListALL == 'binMean_g1')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binKur')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'binSkew')[0][0]
#featureIndex = np.where(statsCollapseListALL == 'multibin_20_ratio_g1')[0][0]



#featureIndex = 12

plt.figure()
plt.grid()
#
### all lead, all no lead
#cut = statsMatrix[:,-1] == 0
#plt.plot(countMatrix[cut], statsMatrix[cut,featureIndex], 'sr', markersize = 5, alpha = 0.3, label = 'Not Pb, Set %d' %setNum)
#cut = statsMatrix[:,-1] == 1
#plt.plot(countMatrix[cut], statsMatrix[cut,featureIndex], 'sb', markersize = 5, alpha = 0.3, label = 'Pb, Set %d' %setNum)
#
# #set 2a
#cut = statsMatrix_0[:,-1] == 0
#plt.plot(countMatrix_0[cut], statsMatrix_0[cut,featureIndex], 'sr', markersize = 5, alpha = 0.3, label = 'Not Pb, Set 2a')
#cut = statsMatrix_0[:,-1] == 1
#plt.plot(countMatrix_0[cut], statsMatrix_0[cut,featureIndex], 'sb', markersize = 5, alpha = 0.3, label = 'Pb, Set 2a')

## set 2b
#cut = statsMatrix_1[:,-1] == 0
#plt.plot(countMatrix_1[cut], statsMatrix_1[cut,featureIndex], 'sr', markersize = 5, alpha = 0.3, label = 'Not Pb, Set 2b')
#cut = statsMatrix_1[:,-1] == 1
#plt.plot(countMatrix_1[cut], statsMatrix_1[cut,featureIndex], 'sb', markersize = 5, alpha = 0.3, label = 'Pb, Set 2b')

#
# #set 2a al and fe, set 2 lead
cut = statsMatrix_0[:,-1] == 0
plt.plot(countMatrix_0[cut], statsMatrix_0[cut,featureIndex], 'sr', markersize = 5, alpha = 0.3, label = 'Not Pb, Set 2a')
cut = statsMatrix[:,-1] == 1
plt.plot(countMatrix[cut], statsMatrix[cut,featureIndex], 'sb', markersize = 5, alpha = 0.3, label = 'Pb, Set 2')

detectorList = goodDetectorsList[(goodDetectorsList > 40) & (goodDetectorsList < 120)]
##groupName = 'PbNOT_1'
groupName = 'Al_0'
#groupName = 'Fe_0'

for i, index in enumerate(datasetGroupsIndices[groupName]):
    plt.plot(stats['count'][index,detectorList], stats[statsCollapseListALL[featureIndex]][index,detectorList], \
        'xm', color = plotColors[i%7], marker = markerTypes[i/7], linestyle = '', \
        markersize = 7, alpha = 0.7, label = '%s, %s'%(filenameList[index], datasetDescription[index]))
#
groupName = 'Fe_0'
for i, index in enumerate(datasetGroupsIndices[groupName]):
    plt.plot(stats['count'][index,detectorList], stats[statsCollapseListALL[featureIndex]][index,detectorList], \
        'xm', color = plotColors[i%7], marker = markerTypes[(i)/7 + 2], linestyle = '', \
        markersize = 7, alpha = 0.7, label = '%s, %s'%(filenameList[index], datasetDescription[index]))

plt.xlabel('Count Per Pulse', fontsize = 16)
plt.ylabel(statsCollapseListALL[featureIndex], fontsize = 16)
plt.legend(loc = 2, prop = {'size':10})

if featureIndex == 1:
    axis((0, 400, 10, 50))
    plt.legend(loc =2, prop = {'size':14})
if featureIndex == 5:
    axis((0, 400, 10, 50))
    plt.legend(loc =1, prop = {'size':10})
elif featureIndex == 9:
    axis((0, 400, 0, 3e4))
elif featureIndex == 11:
    axis((0, 400, 0, 2))
    plt.legend(loc = 3, prop = {'size':10})
elif featureIndex == 12:
    axis((0, 400, 2, 7))
    plt.legend(loc =1, prop = {'size':10})(statsCollapsedMatrix0, statsMatrix0) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed0, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix0, countMatrix0) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed0, ['count'], groupNamesExportList)
countMatrix0 = countMatrix0[:,0]

(statsCollapsedMatrix1, statsMatrix1) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed1, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix1, countMatrix1) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed1, ['count'], groupNamesExportList)
countMatrix1 = countMatrix1[:,0]

(statsCollapsedMatrix2, statsMatrix2) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed2, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix2, countMatrix2) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed2, ['count'], groupNamesExportList)
countMatrix2 = countMatrix2[:,0]

(statsCollapsedMatrix3, statsMatrix3) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed3, statsCollapseListALL, groupNamesExportList)
(countCollapsedMatrix3, countMatrix3) = \
    PBAR_Zspec.CreateStatsCollapsedAll(statsCollapsed3, ['count'], groupNamesExportList)
countMatrix3 = countMatrix3[:,0]
    
plt.axis((0, 400, 3, 7))
plt.axis((0, 400, 10, 50))


# calibration bin value vs detector number

datasetGroupName = 'IB'
datasetGroupName = 'CC'

plotIndexList = datasetGroupsIndices[datasetGroupName]
plt.figure()
plt.grid()

for (dIndex, plotIndex) in enumerate(plotIndexList):
    if dIndex == 0: # skip ec76
        continue
#    plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[plotIndex,goodZspecMask], '.r', color = plotColors[dIndex %7], \
#        linestyle =  lineStyles[dIndex/7], label = '%s, %s' %(filenameList[plotIndex],datasetDescription[plotIndex]))
    plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[plotIndex,goodZspecMask], '.r', color = plotColors[dIndex %7], \
        linestyle =  lineStyles[dIndex/7], label = '%s' %(filenameList[plotIndex]))

# plot the average
plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[:,goodZspecMask].mean(0), '-r', linewidth = 2, label = 'Average')

plt.axis([0, 120, 80, 140])
plt.xlabel('Detector Number', fontsize = 16)
plt.ylabel('Calibration Bin', fontsize = 16)
plt.legend(loc = 3, prop = {'size':8})


#calibration bin value / mean vs detector number

datasetGroupName = 'IB'
datasetGroupName = 'CC'

plotIndexList = datasetGroupsIndices[datasetGroupName]
plt.figure()
plt.grid()

for (dIndex, plotIndex) in enumerate(plotIndexList):
    if dIndex == 0: # skip ec76
        continue
#    plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[plotIndex,goodZspecMask], '.r', color = plotColors[dIndex %7], \
#        linestyle =  lineStyles[dIndex/7], label = '%s, %s' %(filenameList[plotIndex],datasetDescription[plotIndex]))
    plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[plotIndex,goodZspecMask] /gainExtrapolated[:,goodZspecMask].mean(0), '.r', color = plotColors[dIndex %7], \
        linestyle =  lineStyles[dIndex/7], label = '%s' %(filenameList[plotIndex]))

# plot the average
#plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[:,goodZspecMask].mean(0), '-r', linewidth = 2, label = 'Average')

plt.axis([0, 120, 0.9, 1.2])
plt.xlabel('Detector Number', fontsize = 16)
plt.ylabel('Calibration Bin Ratio', fontsize = 16)
plt.legend(loc = 3, prop = {'size':8})


# std/mean of the calibration bin ratio
datasetGroupName = 'IB'
datasetGroupName = 'CC'

plotIndexList = datasetGroupsIndices[datasetGroupName]
plt.figure()
plt.grid()

plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[:,goodZspecMask].std(0)/ gainExtrapolated[:,goodZspecMask].mean(0), \
    '.k', label = '%s' %(filenameList[plotIndex]))

# plot the average
#plt.plot(np.arange(1,137)[goodZspecMask], gainExtrapolated[:,goodZspecMask].mean(0), '-r', linewidth = 2, label = 'Average')

plt.axis([0, 120, 0, 0.1])
plt.xlabel('Detector Number', fontsize = 16)
plt.ylabel('STD/Mean of Calibration Bin', fontsize = 16)


# spectra plots
plotIndexList = datasetGroupsIndices[datasetGroupName]
plotIndexList = [21, 76]
detNum = 101 # starts at 1
temp = [137.5/124.0, 1.0]

plt.figure()
plt.grid()
for (dIndex, datIndex) in enumerate(plotIndexList):
    plt.plot(np.arange(256) * temp[dIndex], dat[datIndex][:,detNum-1] / datasetAcquisitionTime[datIndex] / pulseRate[datIndex] / temp[dIndex], color = plotColors[dIndex %7], \
        linestyle =  lineStyles[dIndex/7], label = '%s, Det# %d, %s' %(filenameList[datIndex], detNum, datasetDescription[datIndex]))
plt.xlabel('Bin', fontsize = 16)
plt.ylabel('Counts Per Pulse', fontsize = 16)
plt.legend()
plt.yscale('log')